# -*- coding: utf-8 -*-
import os
import sys
import socket
import time
import json
from threading import Thread, Lock
import selectors
import types
import math
import struct
import subprocess

import rospy
from nav_msgs.msg import OccupancyGrid, Odometry
from sensor_msgs.msg import BatteryState, Imu
from std_msgs.msg import Float32, String, Bool
from geometry_msgs.msg import Twist
from mirobot.srv import SetGcodeCmd, SetGcodeCmdRequest, GetPoseCmd, GetPoseCmdRequest 



sel = selectors.DefaultSelector()

class ROS_Server:

    def __init__(self) -> None: 
        self.init_tcp()
        self.threads = []
        self.data_lock = Lock()

        # ROS数据变量
        self.map_data = None
        self.odom_data = None
        self.battery_state = None
        self.speed = 0.0
        self.angle = 0.0
        self.mirobot_status = "disconnected"
        self.astra_status = False
        self.emergency = False
        self.control_mode = "manual"
        self.imu_data = {}
        self.client_socket = None
        self.connections = [] 

        self.current_conn_data = None  # 新增当前连接数据对象
        self.outb_lock = Lock()       # 新增输出缓冲区锁
        self.shutdown_flag = False  # 添加关闭标志

        # 初始化ROS节点
        rospy.init_node('tcp_bridge')
        self.cmd_vel_pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
        self.twist_msg = Twist()
        self.linear_speed = 0.2
        self.angular_speed = 0.5

        self.map_sub = rospy.Subscriber("/map", OccupancyGrid, self.map_callback)
        self.current_map = None

        # 初始化服务代理
        self.service_proxies = {}
        self.init_getPoseCmdServices()
        self.init_setGcodeCmdServices()
        
        # 订阅ROS话题
        rospy.Subscriber('/odom', Odometry, self.odom_callback)
        rospy.Subscriber('/battery', BatteryState, self.battery_callback)
        rospy.Subscriber('/speed', Float32, self.speed_callback)
        # rospy.Subscriber('/kata/joints_angle', Float32, self.angle_callback)
        # rospy.Subscriber('/kata/status', String, self.mirobot_callback)
        # rospy.Subscriber('/astra/status', Bool, self.astra_callback)
        rospy.Subscriber('/emergency_stop', Bool, self.emergency_callback)
        # rospy.Subscriber('/control_mode', String, self.control_mode_callback)
        rospy.Subscriber('/imu/data', Imu, self.imu_callback)

    def init_getPoseCmdServices(self):
        
        # 机械臂状态变量初始化
        self.robotCurrentState = {
            'result': 0,
            'state': 0,
            'position': [0.0] * 6,  # x,y,z,a,b,c
            'joint_angles': [0.0] * 7  # jointAngle_1~7
        }
        
        try:
            rospy.wait_for_service('/MirobotServer/GetPoseCmd', timeout=5.0)
            self.get_pose_service = rospy.ServiceProxy('/MirobotServer/GetPoseCmd', GetPoseCmd)
            rospy.loginfo(f"GetPoseCmd Service is ready")
        
        except Exception as e:
            rospy.loginfo(f"service not available will retry later")
        
        # 设置定时器（每100ms请求一次，频率按需调整）
        self.update_timer = rospy.Timer(rospy.Duration(0.5), self.update_robot_status_callback) 

    def update_robot_status_callback(self, event):
        """定时器回调：请求服务并更新状态"""
        try:
            # 创建空请求（该服务可能无需请求参数）
            req = GetPoseCmdRequest()
            # 调用服务
            resp = self.get_pose_service(req)
            # resp = self.get_pose_service.call(req, timeout=rospy.Duration(0.1))
            # rospy.loginfo(f"response {resp.state}")
            
            # 更新状态数据
            self.robotCurrentState['result'] = resp.result
            self.robotCurrentState['state'] = resp.state
            self.robotCurrentState['position'] = [resp.x, resp.y, resp.z, resp.a, resp.b, resp.c]
            self.robotCurrentState['joint_angles'] = [
                resp.jointAngle_1, resp.jointAngle_2, resp.jointAngle_3,
                resp.jointAngle_4, resp.jointAngle_5, resp.jointAngle_6,
                resp.jointAngle_7
            ]
            
        except rospy.ServiceException as e:
            rospy.logerr("Service call failed: %s", str(e))

    def init_setGcodeCmdServices(self):
        """初始化所有需要的服务代理"""

        # GCode 服务
        self.service_proxies['/MirobotServer/SetGcodeCmd'] = rospy.ServiceProxy(
            '/MirobotServer/SetGcodeCmd', SetGcodeCmd)

        for service_name in self.service_proxies:
            try:
                rospy.wait_for_service(service_name, timeout=5.0)
                rospy.loginfo(f"Service {service_name} is ready")
            
            except rospy.ROSException:
                rospy.logwarn(f"Service {service_name} not available, will retry later")

    # ROS回调函数
    def map_callback(self, msg):
        """处理地图数据回调"""
        try:
            self.current_map = msg
            # 如果TCP服务器已连接客户端，立即发送地图数据
            if self.client_socket:
                send_map_thread = Thread(target= self.send_map_data())
                send_map_thread.start()
        
        except Exception as e:
            rospy.logerr(f"Map callback error: {str(e)}")

    def send_map_data(self):
        """发送地图数据给所有连接的客户端"""
        if not self.current_map:
            return
            
        try:
            # 打包地图头信息
            header = struct.pack(
                '!IIddd', 
                self.current_map.info.width,
                self.current_map.info.height,
                self.current_map.info.resolution,
                self.current_map.info.origin.position.x,
                self.current_map.info.origin.position.y
            )
            
            # 处理地图数据：将int8数组转换为bytearray
            # 原始数据中：-1=未知，0=空闲，100=障碍物
            # 转换为：128=未知，255=空闲，0=障碍物
            map_data = bytearray()
            for value in self.current_map.data:
                if value == -1:    # 未知区域
                    map_data.append(128)
                elif value == 0:   # 自由空间
                    map_data.append(255)
                else:              # 障碍物
                    map_data.append(0)
            
            # 组合完整数据包
            data_to_send = b'MAP' + header + map_data
            
            # 发送给所有客户端
            self.client_socket.sendall(data_to_send)
            # for client in self.client_socket:
            #     try:
            #         client.send(data_to_send)
            #     except Exception as e:
            #         rospy.logwarn(f"Failed to send map data to client: {str(e)}")
            #         # 移除断开连接的客户端
            #         self.client_socket.remove(client)
                    
        except Exception as e:
            rospy.logerr(f"Error processing map data: {str(e)}")

    def odom_callback(self, msg):
        with self.data_lock:
            x = msg.pose.pose.position.x
            y = msg.pose.pose.position.y
            q = msg.pose.pose.orientation
            theta = math.atan2(2*(q.w*q.z + q.x*q.y), 1-2*(q.y**2 + q.z**2))
            linear_speed = msg.twist.twist.linear.x
            angular_speed = msg.twist.twist.angular.z
            self.odom_data = {
                'x': x,
                'y': y,
                'theta': theta,
                'linear_speed': linear_speed,
                'angular_speed': angular_speed
            }

    def battery_callback(self, msg):
        with self.data_lock:
            self.battery_state = msg.percentage

    def speed_callback(self, msg):
        with self.data_lock:
            self.speed = msg.data

    def angle_callback(self, msg):
        with self.data_lock:
            self.angle = msg.data

    def mirobot_callback(self, msg):
        with self.data_lock:
            self.mirobot_status = msg.data

    def astra_callback(self, msg):
        with self.data_lock:
            self.astra_status = msg.data

    def emergency_callback(self, msg):
        with self.data_lock:
            self.emergency = msg.data

    def control_mode_callback(self, msg):
        pass
        with self.data_lock:
            self.control_mode = msg.data

    def imu_callback(self, msg):
        with self.data_lock:
            self.imu_data = {
                'linear_acceleration': {
                    'x': msg.linear_acceleration.x,
                    'y': msg.linear_acceleration.y,
                    'z': msg.linear_acceleration.z
                },
                'angular_velocity': {
                    'x': msg.angular_velocity.x,
                    'y': msg.angular_velocity.y,
                    'z': msg.angular_velocity.z
                }
            }

    def shutdown(self):
        """安全关闭所有资源"""
        self.shutdown_flag = True
        sel.close()
        self.tcp_socket.close()
        rospy.loginfo("Server shutdown complete")

    # 网络通信部分
    def init_tcp(self):
        self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.tcp_socket.bind(('0.0.0.0', 8888))
        self.tcp_socket.listen(10)
        self.tcp_socket.setblocking(False)
        sel.register(self.tcp_socket, selectors.EVENT_READ, data=None)

    def run(self):  
        try:  
            send_thread = Thread(target=self.send_data_loop)
            send_thread.start()
            
            while True:  
                events = sel.select(timeout=1)
                for key, mask in events:
                    if key.data is None:
                        self.accept_wrapper(key.fileobj)
                    else:
                        self.service_connection(key, mask)
        
        except KeyboardInterrupt:  
            self.tcp_socket.close()
            sel.close()
            rospy.signal_shutdown("Keyboard interrupt") 
        except Exception as e:
            rospy.logerr(f"Unexpected error: {str(e)}")
        # finally:
        #     self.shutdown()  # 确保资源清理

    def send_data_loop(self):
        while not rospy.is_shutdown():
            try:
                self.send_data_process()
                time.sleep(0.1)
            
            except Exception as e:
                rospy.logerr(f"Send data error: {e}")

    def send_data_process(self):
        try:
            with self.data_lock:
                message_body = self.create_message_body()
            
            if message_body and self.client_socket : ##and self.current_conn_data
                packet = self.build_packet(message_body)
                # rospy.loginfo(f"packet {packet}")
                # with self.outb_lock:
                #     self.client_socket.sendall(packet)

                # 发送给所有活动连接
                for conn in self.connections[:]:  # 使用副本遍历
                    try:
                        if conn.fileno() != -1:  # 检查连接是否有效
                            data = sel.get_key(conn).data
                            with self.outb_lock:
                                data.outb += packet
                    except (ValueError, KeyError, OSError):
                        # 连接已关闭或无效，从列表中移除
                        self.connections.remove(conn)
                        if self.client_socket == conn:
                            self.client_socket = None
                            self.current_conn_data = None
        
        except Exception as e:
            rospy.logwarn(f"send_data_process Failed to send data: {e}")

    def create_message_body(self):
        return json.dumps({
            "ZBot": "ZBot3",
            "sensorData_": {
                "battery": self.battery_state or 0.0,
                "speed": self.speed,
                "angle": self.robotCurrentState['joint_angles'],
                "mirobotstate": self.robotCurrentState['state'],
                "astra": self.astra_status,
                "emergency": self.emergency,
                "controlMode": self.control_mode
            },
            "imuData": self.imu_data,
            "odomData": self.odom_data or {},
            "batteryData": {
                "percentage": self.battery_state or 0.0,
                "voltage": 24.0  
            }
        }) + '\n'

    def build_packet(self, message_body):
        header = json.dumps({"data_size": len(message_body), "type": "ros"}).encode()
        return struct.pack('I', len(header)) + header + message_body.encode()

    def accept_wrapper(self, sock):
        conn, addr = sock.accept()
        print(f"Connected by {addr}")
        # conn.setblocking(False)
        data = types.SimpleNamespace(addr=addr, inb=b'', outb=b'')
        events = selectors.EVENT_READ | selectors.EVENT_WRITE
        sel.register(conn, events, data=data)
        self.connections.append(conn)
        self.client_socket = conn
        self.current_conn_data = data  # 保存当前连接的data对象

    def close_connection(self, sock, data):
        """安全关闭连接并清理资源"""
        try:
            print(f"Closing connection to {data.addr}")
            sel.unregister(sock)
            sock.close()
            
            # 从连接列表中移除
            if sock in self.connections:
                self.connections.remove(sock)
                
            # 如果这是当前活动连接，则重置状态
            if self.client_socket == sock:
                self.client_socket = None
                self.current_conn_data = None
                
        except Exception as e:
            print(f"Error closing connection: {str(e)}")

    def service_connection(self, key, mask):
        sock = key.fileobj
        data = key.data
        if mask & selectors.EVENT_READ:
            try:
                recv_data = sock.recv(4096)
                if recv_data:
                    self.parse_data(recv_data.decode('utf-8'))
                else:
                    self.close_connection(sock, data)
            except:
                print(f"Connection error with {data.addr}")
                self.close_connection(sock, data)

        # if mask & selectors.EVENT_WRITE and data.outb:
        #     try:
        #         sent = sock.send(data.outb)
        #         data.outb = data.outb[sent:]
        #     except:
        #         print("Failed to send data")
        if mask & selectors.EVENT_WRITE:
            with self.outb_lock:
                if data.outb:
                    try:
                        sent = sock.send(data.outb)
                        data.outb = data.outb[sent:]
                    except BlockingIOError:
                        pass
                    except Exception as e:
                        print(f"service_connection -- Failed to send data: {e}")
                        self.close_connection(sock, data)

    def parse_data(self, recv_data):
        try:
            rospy.loginfo(f"收到 JSON 消息: {recv_data}")
            # recv_data = json.loads(recv_data)
            if "MOVE" in recv_data:
                self.handle_rosCarDirectcommand(recv_data)
                rospy.loginfo(f"1111")           
            else:
                index = recv_data.find('}')
                if index != -1:
                    first_part = recv_data[:index+1]
                    second_part = recv_data[index+1:]

                # if isinstance(second_part, str):  
                    # recv_str = json.loads(second_part)
                recv_str = json.loads(second_part)
                # print(f"2025--{recv_str}")
                if 'cmdType' in recv_str: 
                    commandType = recv_str["cmdType"]
                    # commandData = recv_str["cmdData"]
                    commandData = recv_str.get("cmdData")
                    commandBag = recv_str.get("rosBag")

                    if commandType.get('roscore') == "roscore":
                        subprocess.Popen(['xfce4-terminal', '--command', 'bash -c "roscore; exec bash"', '--hold'])
                    
                    if commandType.get('rosrun') == "rosrun":
                        subprocess.Popen(['xfce4-terminal', '--command', 'bash -c "rosrun mirobot MirobotServer /dev/ttyUSB0"', '--hold'])
                        time.sleep(2)
                    elif commandType.get('rosrun') == "rviz":
                        # subprocess.Popen(['xfce4-terminal', '--command', 'bash -c "rviz; exec bash"', '--hold'])
                        terminal_command = ['xfce4-terminal', '--command', f'bash -c "roslaunch zbot3_drive zbot3_slam.launch; exec bash"', '--hold']
                        subprocess.Popen(terminal_command)
                    elif commandType.get('rosrun') == "saveMap":
                        terminal_command = ['xfce4-terminal', '--command', f'bash -c "timeout 10 roslaunch zbot3_drive zbot3_savemap.launch && exit"']
                        subprocess.Popen(terminal_command)
                    elif commandType.get('rosrun') == "carNavConfig":
                        terminal_command = ['xfce4-terminal', '--command', f'bash -c "timeout 10 roslaunch zbot3_drive zbot3_movebase.launch && exit"']
                        subprocess.Popen(terminal_command)
                        time.sleep(1)
                        subprocess.Popen(['xfce4-terminal', '--command', 'bash -c "rosrun rqt_gui rqt_gui"', '--hold'])
                    elif commandType.get('rosrun') == "pointA":
                        terminal_command = ['xfce4-terminal', '--command', f'bash -c "roslaunch zbot3_drive zbot3_navigation.launch; exec bash"', '--hold']
                        # terminal_command = ['xfce4-terminal', '--command', f'bash -c "timeout 10 roslaunch zbot3_drive zbot3_navigation.launch && exit"']
                        subprocess.Popen(terminal_command)

                    if commandType.get('roslaunch') == "roslaunch":
                        if commandData is not None and commandBag is not None:
                            terminal_command = ['xfce4-terminal', '--command', f'bash -c "roslaunch {commandBag} {commandData}; exec bash"', '--hold']
                            subprocess.Popen(terminal_command)
                    
                    if commandType.get('topic') == "topic":
                        subprocess.Popen(['xfce4-terminal', '--command', 'bash -c "rostopic echo /zbot/odom; exec bash"', '--hold'])
                    
                    if commandType.get('service') == "service":
                        if commandData is not None:
                            terminal_command = ['xfce4-terminal', '--command', f'bash -c "timeout 10 rosservice call {commandData} && exit"']
                            subprocess.Popen(terminal_command)
                        else:
                            subprocess.Popen(['xfce4-terminal', '--command', 'bash -c "timeout 10 rosservice call /MirobotServer/SetHomeCmd && exit"'])

                    if commandType.get('GCode') == "GCode":
                        self.send_gcode_command(commandData)

        except json.JSONDecodeError as e:  
            print(f"JSON解析错误: {e}")  

    def send_gcode_command(self, index):
        try:
            # 映射索引到 GCode 字符串
            gcode_map = {
                1: "M21 G91 G1 X10 F1000",
                2: "M21 G91 G1 X-10 F1000",
                3: "M21 G91 G1 Y10 F1000",
                4: "M21 G91 G1 Y-10 F1000",
                5: "M21 G91 G1 Z10 F1000",
                6: "M21 G91 G1 Z-10 F1000",
                7: "M21 G91 G1 A10 F1000",
                8: "M21 G91 G1 A-10 F1000",
                9: "M21 G91 G1 B10 F1000",
                10: "M21 G91 G1 B-10 F1000",
                11: "M21 G91 G1 C10 F1000",
                12: "M21 G91 G1 C-10 F1000",
                21: "o117",
                22: "$h",  # Home 命令
                23: "M21 G90 G00 X0 Y0 Z0 A0 B0 C0 F2000"
            }
            gcode_str = gcode_map.get(int(index))
            rospy.loginfo(f"GCode 1")
            if not gcode_str:
                rospy.logwarn(f"Invalid GCode index: {index}")
                return
                
            try:
                rospy.loginfo(f"GCode 2")
                # 确保服务代理存在
                service_name = '/MirobotServer/SetGcodeCmd'
                if service_name not in self.service_proxies:
                    self.service_proxies[service_name] = rospy.ServiceProxy(service_name, SetGcodeCmd)
                rospy.loginfo(f"GCode 3")
                # 等待服务可用
                rospy.wait_for_service(service_name, timeout=2.0)
                rospy.loginfo(f"GCode 4")
                # 创建请求对象
                req = SetGcodeCmdRequest()
                req.gcode = gcode_str
                rospy.loginfo(f"GCode 5")
                # # 调用服务
                response = self.service_proxies[service_name](req)
                # rospy.loginfo(f"GCode command sent: {gcode_str}, success: {response.success}")
                
                # # 调用服务 - 增强错误处理
                # try:
                #     service_response = self.service_proxies[service_name](req)
                #     success = service_response.success
                #     message = "GCode 命令已执行"
                #     rospy.loginfo(f"GCode 命令成功: {gcode_str}")
                # except rospy.ServiceException as se:
                #     success = False
                #     message = f"服务调用失败: {str(se)}"
                #     rospy.logerr(message)
                # except rospy.ROSException as re:
                #     success = False
                #     message = f"服务不可用: {str(re)}"
                #     rospy.logerr(message)
                # except Exception as e:
                #     success = False
                #     message = f"意外错误: {str(e)}"
                #     rospy.logerr(message)

                
            except rospy.ServiceException as e:
                rospy.logerr(f"GCode service call failed: {str(e)}")
            except rospy.ROSException as e:
                rospy.logerr(f"GCode service unavailable: {str(e)}")

        except rospy.ServiceException as e:
            rospy.logerr(f"Service call failed: {e}")

    def handle_rosCarDirectcommand(self, command):
        rospy.loginfo(f"2222")
        # 命令格式: "MOVE:<char>:<linear>:<angular>:<omni>"
        parts = command.split(':')
        if len(parts) < 5 or parts[0] != "MOVE":
            return False
        
        cmd_char = parts[1]
        
        try:
            linear = float(parts[2])
            angular = float(parts[3])
            omni = int(parts[4]) == 1
        
        except:
            return False
        
        ## 更新速度参数
        self.linear_speed = linear
        self.angular_speed = angular
        self.is_omni_mode = omni
        
        ## 根据命令字符设置速度
        self.twist_msg.linear.x = 0.0
        self.twist_msg.linear.y = 0.0
        self.twist_msg.angular.z = 0.0
        
        if cmd_char == 'u':  # 左上
            self.twist_msg.linear.x = self.linear_speed
            self.twist_msg.angular.z = self.angular_speed
        elif cmd_char == 'i':  # 上
            self.twist_msg.linear.x = self.linear_speed
        elif cmd_char == 'o':  # 右上
            self.twist_msg.linear.x = self.linear_speed
            self.twist_msg.angular.z = -self.angular_speed
        elif cmd_char == 'j':  # 左
            self.twist_msg.angular.z = self.angular_speed
        elif cmd_char == 'l':  # 右
            self.twist_msg.angular.z = -self.angular_speed
        elif cmd_char == 'm':  # 左下
            self.twist_msg.linear.x = -self.linear_speed
            self.twist_msg.angular.z = self.angular_speed
        elif cmd_char == ',':  # 下
            self.twist_msg.linear.x = -self.linear_speed
        elif cmd_char == '.':  # 右下
            self.twist_msg.linear.x = -self.linear_speed
            self.twist_msg.angular.z = -self.angular_speed
        # 'k' 停止，不做任何操作
        
        ## 全向轮模式特殊处理
        if self.is_omni_mode:
            # 这里需要根据全向轮模型调整速度分量
            # 示例：将角速度转换为横向分量
            self.twist_msg.linear.y = self.twist_msg.angular.z
            self.twist_msg.angular.z = 0.0
        rospy.loginfo(f"333")
        ## 发布速度命令
        self.cmd_vel_pub.publish(self.twist_msg)
        rospy.loginfo(f"44")
        return True


if __name__ == '__main__':
    server = ROS_Server()
    server.run()